import os
import re
import sys
import json
from dataclasses import dataclass, field
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Dict, List, Any, Tuple

import numpy as np
import torch
from jsonargparse import ArgumentParser
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset

from mcap_base_reader import MCAPBaseReader, SyncedMCAPData, DEFAULT_TASK_INFO_PATH
from data_asst.config.robot_config import ROBOT_MAP, RobotFeaturesBase


"""
使用示例：
python data_asst/mcap_utils/mcap_2_lerobot.py \
    --task_name "pick_bottle_table_0927 pick_bread_0927" \
    --input_dir "/path/to/mcap_data" \
    --output_dir "/path/to/lerobot_datasets" \
    --robot_config "airbot_dual"
"""


@dataclass(frozen=True)
class ConversionConfig:
    """批量将 MCAP 文件转换为 LeRobotDataset 格式的配置"""

    task_name: str = field(default="")
    input_dir: str = field(default="")
    output_dir: str = field(default="")
    robot_config: str = field(default="airbot_dual")
    max_workers: int = field(default=max(1, os.cpu_count() // 2))

    def __post_init__(self):
        object.__setattr__(self, "input_dir", Path(self.input_dir))
        object.__setattr__(self, "output_dir", Path(self.output_dir))

        # 拆分 task_name
        task_names_list = []
        if isinstance(self.task_name, str):
            task_names_list = self.task_name.split()
            # 同样更新 self.task_name，因为 main() 会用到它
            object.__setattr__(self, "task_name", task_names_list)
        else:
            task_names_list = self.task_name

        if not task_names_list:
            # 如果没有任务，也动态创建空列表，防止 main() 校验时属性不存在
            object.__setattr__(self, "task_prompt", [])
            object.__setattr__(self, "ignore_episodes", [])
            return

        # 加载 task_info.json 并创建一个 map
        task_info_map = {}
        try:
            with open(DEFAULT_TASK_INFO_PATH, "r", encoding="utf-8") as f:
                task_data = json.load(f)
            task_info_map = {item['task_name']: item for item in task_data}
            print(f"[*] 成功从 {DEFAULT_TASK_INFO_PATH} 加载了 {len(task_info_map)} 条任务信息。")
        except Exception as e:
            print(f"[!] 错误: 无法加载或解析 task_info.json (路径: {DEFAULT_TASK_INFO_PATH})。错误: {e}")
            print("[!] 程序将退出。")
            # 动态创建空列表，以便 main() 中的校验失败
            object.__setattr__(self, "task_prompt", [f"Error loading prompt for {t}" for t in task_names_list])
            object.__setattr__(self, "ignore_episodes", [[] for _ in task_names_list])
            # 在这种情况下，后续的校验应该会失败并退出
            return

        new_prompts = []
        new_ignores = []

        # 遍历 task_name 并从 JSON 填充
        for t_name in task_names_list:
            if t_name in task_info_map:
                new_prompts.append(task_info_map[t_name]['task_instruction'])
                new_ignores.append(task_info_map[t_name]['ignore_episodes'])
            else:
                print(f"[*] 警告: 任务 '{t_name}' 在 {DEFAULT_TASK_INFO_PATH} 中未找到。将使用默认 prompt 和空 ignore list。")
                new_prompts.append(f"Default prompt for task {t_name}")
                new_ignores.append([])
        
        # 动态地将这两个列表添加到配置实例上
        object.__setattr__(self, "task_prompt", new_prompts)
        object.__setattr__(self, "ignore_episodes", new_ignores)


@dataclass(frozen=True)
class DatasetConfig:
    """LeRobotDataset 的创建参数"""
    use_videos: bool = True
    tolerance_s: float = 0.0001
    image_writer_processes: int = 10
    image_writer_threads: int = 5
    video_backend: str | None = None


DEFAULT_DATASET_CONFIG = DatasetConfig()


class MCAPReader(MCAPBaseReader):
    """
    读取并处理单个MCAP文件，将其内容提取为适合LeRobotDataset的帧列表。
    """

    def __init__(self, mcap_file: Path):
        super().__init__(mcap_file)

    def process_mcap_to_frames(self) -> List[Dict[str, Any]]:
        """
        主处理函数：从基类获取数据，并将其转换为 LeRobot 帧列表。
        """
        synced_data = self.read_and_sync_data()

        if synced_data is None:
            print(
                f"[*] 警告: 文件 {self.mcap_file.name} 缺少有效的关节或图像数据，已跳过。"
            )
            return []

        return self._prepare_lerobot_frames(synced_data)

    def _prepare_lerobot_frames(
        self, synced_data: SyncedMCAPData
    ) -> List[Dict[str, Any]]:
        """
        将 SyncedMCAPData 转换为 LeRobotDataset 所需的帧字典列表。
        """
        eff_len = len(synced_data.timestamps)
        if eff_len == 0:
            return []

        all_frames_for_episode = []

        # 预取数据，避免在循环中重复索引
        qpos_obs_all = synced_data.qpos_obs
        qpos_act_all = synced_data.qpos_act
        gpos_obs_all = synced_data.gpos_obs
        gpos_act_all = synced_data.gpos_act

        # images_by_hdf5_cam: {"cam_left": [T,H,W,C], ...}
        images_by_hdf5_cam = {}
        for hdf5_cam, mcap_cam_key in self.CAMERA_MAPPING.items():
            images_by_hdf5_cam[hdf5_cam] = synced_data.images[mcap_cam_key]

        for i in range(eff_len):
            # 准备 qpos 和 action
            qpos = qpos_obs_all[i]
            action_qpos = qpos_act_all[i]
            gpos = gpos_obs_all[i]
            action_gpos = gpos_act_all[i]

            frame_dict = {
                "observation.state": torch.from_numpy(qpos),
                "action": torch.from_numpy(action_qpos),
                "observation.state_gpos": torch.from_numpy(gpos),
                "action_gpos": torch.from_numpy(action_gpos),
            }

            # 添加图像
            for hdf5_cam in self.camera_names:  # self.camera_names = ["cam_left", ...]
                frame_dict[f"observation.images.{hdf5_cam}"] = images_by_hdf5_cam[
                    hdf5_cam
                ][i]

            all_frames_for_episode.append(frame_dict)

        return all_frames_for_episode


def process_single_task(
    task_name: str, task_prompt: str, ignored_indices: set, config: ConversionConfig
):
    """处理单个任务的完整流程：创建数据集、查找文件、转换并保存。"""
    print(f"🚀 开始处理任务: {task_name}")
    task_input_dir = config.input_dir / task_name
    dataset_path = config.output_dir / task_name
    robot_config_obj = ROBOT_MAP[config.robot_config]()

    if not task_input_dir.is_dir():
        print(
            f"[*] 警告: 任务 '{task_name}' 的输入目录 {task_input_dir} 不存在，已跳过。"
        )
        return

    mcap_files = sorted(
        list(task_input_dir.rglob("*.mcap")),
        key=lambda p: (
            int(re.search(r"(\d+)", p.name).group(1))
            if re.search(r"(\d+)", p.name)
            else -1
        ),
    )

    jobs = [
        mcap_file
        for mcap_file in mcap_files
        if int(re.search(r"(\d+)", mcap_file.name).group(1)) not in ignored_indices
    ]

    if not jobs:
        print(f"[*] 信息: 任务 '{task_name}' 没有需要处理的文件。")
        return

    dataset = LeRobotDataset.create(
        repo_id=dataset_path.name,
        root=dataset_path,
        fps=robot_config_obj.fps,
        features=robot_config_obj.features(
            has_velocity=False, has_effort=False, image_dtype="video"
        ),
        use_videos=DEFAULT_DATASET_CONFIG.use_videos,
        tolerance_s=DEFAULT_DATASET_CONFIG.tolerance_s,
        image_writer_processes=DEFAULT_DATASET_CONFIG.image_writer_processes,
        image_writer_threads=DEFAULT_DATASET_CONFIG.image_writer_threads,
        video_backend=DEFAULT_DATASET_CONFIG.video_backend,
    )

    # (注意：这里为了简单起见，在主进程中处理文件。
    #  如果追求极致性能，可以将 MCAPReader 的调用也放入多进程，
    #  然后让主进程的 dataset.add_frame 消费结果。)
    for mcap_file in jobs:
        reader = MCAPReader(mcap_file)  # MCAPReader 现在是重构后的版本
        all_frames = reader.process_mcap_to_frames()
        if not all_frames:
            continue
        for frame in all_frames:
            frame["task"] = task_prompt
            dataset.add_frame(frame)
        dataset.save_episode()

    print(f"✅ 任务 '{task_name}' 处理完成，数据集保存在: {dataset_path}")


def main():
    parser = ArgumentParser()
    parser.add_class_arguments(ConversionConfig, as_group=False)
    cfg = parser.parse_args()
    config = ConversionConfig(**vars(cfg))

    print(
        f"任务: {config.task_name}\n输入目录: {config.input_dir}\n输出目录: {config.output_dir}"
    )
    print(f"--- 自动加载的 Prompts (来自 {DEFAULT_TASK_INFO_PATH}) ---")
    if config.task_prompt:
        for i, p in enumerate(config.task_prompt):
            print(f"  [{config.task_name[i]}]: {p[:50]}...") # 打印前50个字符
    else:
        print("  (未加载任何 task prompts)")
    print(f"--------------------------")

    # --- 参数校验 ---
    if len(config.task_name) != len(config.ignore_episodes):
        print(
            f"[!] 错误: 任务名数量({len(config.task_name)})和忽略列表数量({len(config.ignore_episodes)})必须一致！"
        )
        sys.exit(1)
    if len(config.task_name) != len(config.task_prompt):
        print(
            f"[!] 错误: 任务名数量({len(config.task_name)})和任务描述数量({len(config.task_prompt)})必须一致！"
        )
        sys.exit(1)
    if not config.task_name:
        print("[!] 错误: 未指定任何 task_name。")
        return

    # --- 使用多进程并行处理每个任务 ---
    print(
        f"\n准备使用 {config.max_workers} 个进程并行处理 {len(config.task_name)} 个任务..."
    )
    with ProcessPoolExecutor(max_workers=config.max_workers) as executor:
        # 建立 future 到 task_name 的映射，方便后续打印信息
        future_to_task = {
            executor.submit(
                process_single_task,
                config.task_name[i],
                config.task_prompt[i],
                set(config.ignore_episodes[i]),
                config,
            ): config.task_name[i]
            for i in range(len(config.task_name))
        }

        # 这里不再使用外层tqdm，避免进度条冲突
        completed_count = 0
        total_tasks = len(config.task_name)
        for future in as_completed(future_to_task):
            task_name = future_to_task[future]
            completed_count += 1
            try:
                future.result()  # 获取结果，同时也能捕获子进程中的异常
                print(
                    f"[{completed_count}/{total_tasks}] 任务 '{task_name}' 已成功完成。"
                )
            except Exception as e:
                print(
                    f"[!!!] [{completed_count}/{total_tasks}] 任务 '{task_name}' 执行时发生严重错误: {e}"
                )

    print("\n--- 所有任务转换完成 ---")


if __name__ == "__main__":
    main()
